import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import requests
import tiktoken


class DataHandler:
    # 数据处理
    def __init__(self, context_length=64):
        self.context_length = context_length
        self.vocab_size = 0
        self.train_data = None
        self.valid_data = None
        self.lode_build_date()

    def lode_build_date(self):
        self.building_data(self._load_data())

    @staticmethod
    def _load_data():
        # Download dataset if not exists
        if not os.path.exists('sales_textbook.txt'):
            url = "https://huggingface.co/datasets/goendalf666/sales-textbook_for_convincing_and_selling/resolve/main/sales_textbook.txt?download=true"
            with open('sales_textbook.txt', 'wb') as f:
                f.write(requests.get(url).content)

        # Read and tokenize text
        with open('sales_textbook.txt', 'r', encoding="utf-8") as f:
            text = f.read()
        return text

    def building_data(self, text):
        encoding = tiktoken.get_encoding("cl100k_base")
        tokenized_text = torch.tensor(encoding.encode(text), dtype=torch.long)
        self.vocab_size = len(set(tokenized_text))
        # self.vocab_size = tokenized_text.max().item() + 1

        # Train/validation split
        train_idx = int(len(tokenized_text) * 0.9)
        self.train_data = tokenized_text[:train_idx]
        self.valid_data = tokenized_text[train_idx:]

    def get_batch(self, batch_size, split='train'):
        data = self.train_data if split == 'train' else self.valid_data
        indexes = torch.randint(0, len(data) - self.context_length, (batch_size,))

        x_batch = torch.stack([data[i:i + self.context_length] for i in indexes])
        y_batch = torch.stack([data[i + 1:i + self.context_length + 1] for i in indexes])
        return x_batch, y_batch


class FeedForwardNetwork(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_model * 4)
        self.linear2 = nn.Linear(d_model * 4, d_model)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        x = self.dropout(x)
        return x


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        self.d_model = d_model

        # Create positional encoding table
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # Add batch dimension
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]


# 3. 多头注意力模块
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, context_length):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

        # Create causal mask
        self.register_buffer("mask", torch.triu(torch.ones(context_length, context_length), diagonal=1).bool())

    def forward(self, x):
        batch_size, context_length, _ = x.shape

        # Linear projections
        q_ = self.w_q(x)
        k_ = self.w_k(x)
        v_ = self.w_v(x)

        # Split into multiple heads
        q = q_.view(batch_size, context_length, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        k = k_.view(batch_size, context_length, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        v = v_.view(batch_size, context_length, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        # Scaled dot-product attention
        attention_scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attention_mask_scores = attention_scores.masked_fill(self.mask, float('-inf'))
        attention_weights = F.softmax(attention_mask_scores, dim=-1)

        # Apply attention to values
        output_a = attention_weights @ v

        # Concatenate heads
        output = output_a.permute(0, 2, 1, 3).contiguous()
        output = output.view(batch_size, context_length, self.d_model)

        # Final linear projection
        output_o = self.w_o(output)

        return output_o


class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, context_length):
        super().__init__()
        self.attention = MultiHeadAttention(d_model, num_heads, context_length)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = FeedForwardNetwork(d_model)

    def forward(self, x):
        # Self-attention with residual connection
        attn_output = self.attention(x)
        x = self.norm1(x + attn_output)

        # Feed-forward network with residual connection
        ffn_output = self.ffn(x)
        return self.norm2(x + ffn_output)


class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model=512, num_heads=8, num_layers=6):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([TransformerBlock(d_model, num_heads, context_length) for _ in range(num_layers)])
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        # Input embedding with positional encoding
        x = self.embedding(x)
        x = self.pos_encoder(x)

        # Pass through transformer layers
        for layer in self.layers:
            x = layer(x)

        # Final projection to vocabulary size
        return self.fc_out(x)


# Usage example
if __name__ == "__main__":
    # Initialize data handler
    data_handler = DataHandler(context_length=64)

    # Initialize model
    model = Transformer(vocab_size=data_handler.vocab_size)

    # Get a batch
    x_batch, y_batch = data_handler.get_batch(batch_size=32)

    # Forward pass
    logits = model(x_batch)
    print(f"Input shape: {x_batch.shape}")
    print(f"Output shape: {logits.shape}")  # Should be [batch_size, context_length, vocab_size]
